{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# INFO8010: Homework 1\n",
    "\n",
    "In this first homework, you will familiarize yourself with the basics of differentiable programming and the PyTorch framework. At the end of this assignment, you should be able to build and train your own multi-layer perceptron. Although this assignment is optional, we **strongly** advise you to do it as you are expected to be comfortable with PyTorch at the end of this course.\n",
    "\n",
    "In the following, we provide some code blocks that are already running, along with explanations. You are asked to complete the code blocks in which the `# your code` comment is present. Sometimes, you will also have to motivate why you programmed certain things or guide us quickly through the results you obtained. If this is required you will see the following instruction:\n",
    "\n",
    "> your discussion\n",
    "\n",
    "Please note that you will **not** have to handle in any sort of written report by the end of the assignment. We do however expect you to submit the notebook **with the solutions** to the exercises."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Imports\n",
    "\n",
    "Our first step is very easy: we simply import the required libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Tensors and basic operations\n",
    "\n",
    "Tensors are one of the main ingredients when it comes to modern deep learning frameworks. Almost all deep learning computations can be expressed as tensor operations which make computation fast and efficient, especially on graphics processing units (GPUs). We will now see how to manipulate tensors within the PyTorch framework.\n",
    "\n",
    "There are many valid definitions for tensors, but to keep it simple you can see them as multi-dimensional arrays. Zero-dimensional tensors are scalars, one-dimensional tensors are vectors, two-dimensional tensors are matrices, etc.\n",
    "\n",
    "In PyTorch, a tensor is an instance of the class `torch.Tensor`. PyTorch tensors are **very** similar to NumPy arrays.\n",
    "Almost any operations you could imagine can be performed on tensors (see [API](https://pytorch.org/docs/stable/tensors.html))."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Tensors\n",
    "\n",
    "There are many ways to [create tensors](https://pytorch.org/docs/stable/tensors.html#torch.Tensor), for instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we can create a tensor from lists\n",
    "t1 = torch.tensor([[1, 2], [-7, 9]])\n",
    "print(t1)\n",
    "\n",
    "# or from numpy arrays\n",
    "t2 = torch.from_numpy(np.array([[1., 8.], [0., 3.]]))\n",
    "print(t2)\n",
    "\n",
    "# or we can create a 2x3x4 tensor filled with zeros\n",
    "t3 = torch.zeros(2, 3, 4)\n",
    "print(t3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an array may contain different types of data (e.g. float, int, boolean), so can a tensor. Let's check what types of data are in `t1`, `t2` and `t3`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(t1.dtype, t2.dtype, t3.dtype, sep='\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, PyTorch uses an integer data type to store integer values and a floating-point data type to store real values. However, unlike NumPy, PyTorch uses a **single-precision** floating-point data type (`float32`) by default.\n",
    "\n",
    "You can convert tensors from one type to another:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(t1.float(), t2.long(), t3.double(), sep='\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another important property of a tensor is its shape. You can use the `.shape` property or `.size()` method to check the dimensions of a tensor. Both return a `torch.Size` object that can be manipulated with its own operations. Most of the time, you will just use `torch.Size` objects as tuples or lists (e.g. to check the size of a tensor along one of its dimension)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the full shape\n",
    "print(t1.shape)\n",
    "\n",
    "# the length of the tensor along its first dimension\n",
    "print(t2.shape[0])\n",
    "\n",
    "# the shape of the two last dimensions\n",
    "print(t3.shape[-2:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Operations\n",
    "\n",
    "We now consider simple tensor operations on 1D (vectors) and 2D (matrices) tensors. Lets instantiate two vectors with 5 elements and two matrices of respectively 3x5 and 5x5 elements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v1 = torch.ones(5)\n",
    "print(v1)\n",
    "\n",
    "v2 = torch.randn(5)  # standard gaussian\n",
    "print(v2)\n",
    "\n",
    "m1 = torch.rand(3, 5)  # [0, 1) uniform\n",
    "print(m1)\n",
    "\n",
    "m2 = torch.eye(5)  # identity matrix\n",
    "print(m2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most basic operations are the addition, subtraction, multiplication and division. As in NumPy, but unlike MATLAB, the symbols `+, -, *, /` all perform **element-wise** operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add v1 and v2\n",
    "v_sum = v1 + v2\n",
    "print(v_sum)\n",
    "\n",
    "# subtract v1 from v2\n",
    "v_sub = v1 - v2\n",
    "print(v_sub)\n",
    "\n",
    "# multiply the elements of v1 and v2\n",
    "v_mul = v1 * v2\n",
    "print(v_mul)\n",
    "\n",
    "# divide the elements of v1 by v2\n",
    "v_div = v1 / v2\n",
    "print(v_div)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When trying to perform an element-wise operation between tensors of different shapes, the tensors are [broadcast](https://numpy.org/doc/stable/user/basics.broadcasting.html) together, if possible. Notably, this allows to add or multiply tensors by scalars."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add a scalar\n",
    "print(v1 + 1)\n",
    "\n",
    "# multiply by scalar\n",
    "print(v2 * 2)\n",
    "\n",
    "# broadcast substraction between vector and matrix\n",
    "print(v2 - m1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another basic operation with tensors is the inner product, denoted by the symbol `@`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# scalar product between v1 and v2\n",
    "print(v1 @ v2)\n",
    "\n",
    "# matrix-vector product between m1 and v2\n",
    "print(m1 @ v2)\n",
    "\n",
    "# matrix-matrix product between m1 and m2\n",
    "print(m1 @ m2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Aggregation\n",
    "\n",
    "Common operations are to aggregate the values of a tensor along a dimension into a single value. For instance, computing the sum of the rows of a matrix . A lot of aggregation methods such as `.sum`, `.prod`, `.mean`, `.max`, ... are available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sum of rows of m1\n",
    "print(m1.sum(dim=1))\n",
    "\n",
    "# mean of columns of m2\n",
    "print(m2.mean(dim=0))\n",
    "\n",
    "# prod of the elements of v2\n",
    "print(v2.prod())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Indexing\n",
    "\n",
    "Sometimes you would like to extract a sub-tensor from the tensor. This operation is called *slicing*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract elements 0 (included) to 2 (not included)\n",
    "print(v1[0:2])\n",
    "\n",
    "# extract all elements but the last\n",
    "print(v2[:-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may also extract a subset of the elements by passing a list of indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract the first, fourth and fifth elements\n",
    "v2[[0, 3, 4]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matrices can also be sliced/indexed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m1[:4, 2:-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Squeeze and Unsqueeze\n",
    "\n",
    "Squeezing removes dimensions of size 1 from tensors. Unsqueezing adds a dimension of size 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# before unsqueeze\n",
    "print(m1.shape)\n",
    "\n",
    "# after unsqueeze\n",
    "m1 = m1.unsqueeze(1)\n",
    "print(m1.shape)\n",
    "\n",
    "# after squeeze\n",
    "m1 = m1.squeeze(1)\n",
    "print(m1.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### View\n",
    "\n",
    "Sometimes, tensors don't have the correct shape. For example, you might want to process 3x32x32 images as vectors of 3072 elements. The `.view` method returns a new tensor with the same data, but of a different shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = torch.rand(10, 3, 32, 32)\n",
    "print(images.shape)\n",
    "\n",
    "images_as_vectors = images.view(10, -1)\n",
    "print(images_as_vectors.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, both tensors share the same underlying data, meaning that modifying one in-place will also modify the other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images[0, :] = torch.zeros_like(images[0])\n",
    "images_as_vectors[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It should be noted that slicing, indexing and squeezing operations actually create views of the tensor instead of making copies of the data. For a complete list of view operations in `torch`, see the [documentation](https://pytorch.org/docs/stable/tensor_view.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Other primitives\n",
    "\n",
    "A **lot** of useful primitives are available in `torch`. Notable examples are [exp](https://pytorch.org/docs/stable/generated/torch.exp.html), [log](https://pytorch.org/docs/stable/generated/torch.log.html), [sqrt](https://pytorch.org/docs/stable/generated/torch.sqrt.html), [ceil](https://pytorch.org/docs/stable/generated/torch.ceil.html), [clamp](https://pytorch.org/docs/stable/generated/torch.expand.html), [sort](https://pytorch.org/docs/stable/generated/torch.sort.html), [argmax](https://pytorch.org/docs/stable/generated/torch.argmax.html), [stack](https://pytorch.org/docs/stable/generated/torch.stack.html), [cat](https://pytorch.org/docs/stable/generated/torch.cat.html), [where](https://pytorch.org/docs/stable/generated/torch.where.html) and so on. Always take a look at the [documentation](https://pytorch.org/docs/stable/index.html) before implementing something. And don't forget: Google is your best friend!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Try it yourself!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a tensor `samples` with $10^{5}$ i.i.d. normally (mean 5 and standard deviation 2) distributed values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = # your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract the first 42 elements of `samples`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract the last 666 elements of `samples`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract all the elements of `samples` that have an odd index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a view of shape 10x1000x10 of `samples`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the mean of `samples` using a `for` loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the mean and standard deviation of `samples` with the appropriate PyTorch operators:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "\n",
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What do you observe when you compare the running time of the two code snippets?\n",
    "\n",
    "> your discussion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute the absolute-value ($L_1$) and Euclidean ($L_2$) norms of `samples`, without loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract the first and last 13 elements of `samples` as two distinct vectors and compute their [outer product](https://en.wikipedia.org/wiki/Outer_product) matrix, without loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract the diagonal of this matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Trace the curve defined by the coordinates\n",
    "\n",
    "$$x(t) = 16 \\sin(t)^3$$\n",
    "\n",
    "$$y(t) = 13 \\cos(t) - 5 \\cos(2t) - 2 \\cos(3t) - \\cos(4t)$$\n",
    "\n",
    "for $t \\in [0, 2\\pi[$.\n",
    "\n",
    "You can use `np.pi` for the value of $\\pi$ and the `matplotlib` package to plot the curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now send the result to your crush ;)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. The autograd package\n",
    "\n",
    "The autograd package is what really makes PyTorch different from other algebraic language/libraries such as MATLAB or NumPy. Whenever you make operations in PyTorch, it will create a computation graph that can later be used to compute derivatives of the output quantities with respect to the input or other intermediate computation steps. \n",
    "\n",
    "The official documentation says\n",
    "\n",
    "> `torch.autograd` provides classes and functions implementing automatic differentiation of arbitrary scalar valued functions. It requires minimal changes to the existing code - you only need to declare tensors for which gradients should be computed with the `requires_grad=True` keyword.\n",
    "\n",
    "For more information about autograd you can check this short [tutorial](https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html) and the [documentation](https://pytorch.org/docs/stable/autograd.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 A guided tour\n",
    "\n",
    "Let's see how autograd may be used to compute the derivative(s) of the following function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return x**2 + 2 * x - 0.5 * torch.sin(x * 3)\n",
    "\n",
    "def f_prime(x):  # f'\n",
    "    return 2 * x + 2 - 1.5 * torch.cos(x * 3)\n",
    "\n",
    "def f_pprime(x):  # f''\n",
    "    return 2 + 4.5 * torch.sin(x * 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.randn(1, requires_grad=True)  # without requires_grad, the computation graph is not created\n",
    "print(x)\n",
    "\n",
    "y = f(x)\n",
    "print(y)\n",
    "\n",
    "dy = torch.autograd.grad(y, x, create_graph=True)[0]  # create_graph is necessary for higher-order derivative\n",
    "print(dy, f_prime(x))\n",
    "\n",
    "ddy = torch.autograd.grad(dy, x)[0]\n",
    "print(ddy, f_pprime(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, autograd is able to compute gradients automatically. We can exploit this to perform **gradient descent** with parametric approximators. Let's approximate the function `f` using a polynomial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# learning set\n",
    "x = torch.linspace(-3, 3, 1000)\n",
    "y = f(x)\n",
    "\n",
    "# parametric approximator (polynomial of degree 3)\n",
    "def approximator(w, x):\n",
    "    return torch.stack((torch.ones_like(x), x, x**2, x**3), dim=-1) @ w\n",
    "\n",
    "# initialize parameters at random\n",
    "w = torch.randn(4, requires_grad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look at what happens before training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.title('Before training')\n",
    "plt.plot(x, y, label='Ground truth')\n",
    "plt.plot(x, approximator(w, x).detach(), label='Approximation')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We perform gradient descent on the mean squared error (MSE) loss function in order to fit the learning set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = .001  # learning rate\n",
    "\n",
    "for i in range(1000):\n",
    "    y_pred = approximator(w, x)\n",
    "    mse = torch.mean((y - y_pred)**2)  # mean squared error\n",
    "    grad = torch.autograd.grad(mse, w)[0]\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        w -= lr * grad  # gradient descent update\n",
    "\n",
    "plt.title('After training')\n",
    "plt.plot(x, y, label='Ground truth')\n",
    "plt.plot(x, approximator(w, x).detach(), label='Approximation')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the approximator became larger, calling the `torch.autograd.grad` method with all parameter tensors would become quite cumbersome. Fortunately, PyTorch provides the `.backward()` method to acumulate (in the `.grad` property) the gradient values in **all** tensors that are in the computation graph. Later on in this notebook we will see how this, together with `Optimizer` objects, makes gradient descent very simple."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.rand(5)\n",
    "w = torch.rand(5, requires_grad=True)\n",
    "y = w @ x\n",
    "y.backward()\n",
    "\n",
    "print(x, w.grad, sep='\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importantly, the construction of the computation graphs is expensive and memory consuming. If some computations do not use the gradients, like evaluating an approximator or updating the parameters, one can disable the creation of the computation graph with the `torch.no_grad()` context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = torch.rand(100, 100, requires_grad=True)\n",
    "x = torch.rand(100)\n",
    "\n",
    "%timeit w @ x\n",
    "\n",
    "with torch.no_grad():\n",
    "    %timeit w @ x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Try it yourself!\n",
    "\n",
    "We would like you to perform the same kind of curve fitting that we did a few cells above, but this time we want to use the `.backward()` method instead of `torch.autograd.grad`. Remember that PyTorch will keep **accumulating** the gradient values in the `.grad` property of each tensor, i.e. the values are not replaced but added together after each backward pass. Fortunately, you can reset the values of the gradient to 0 with the `.zero_()` in-place method (e.g. `w.grad.zero_()`). To make your life easier, you can start from the learning loop three cells above, but be careful to re-initialize the parameters at random before training and to not use the `torch.autograd.grad` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. The `torch.nn` and `torch.optim` modules\n",
    "\n",
    "You should now be able to define your own neural network and train it with tensorial operations and the autograd package. However this would require you to explicitely define every operation in the neural network and to keep track of all the parameters for performing gradient descent. Fortunately, PyTorch provides the `torch.nn` and `torch.optim` packages which implement everything you need to define and train a neural network efficiently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### `nn.functional`\n",
    "\n",
    "The `torch.nn.functional` (imported as `F` in this notebook) module implements many predefined functions which should simplify your life for building your own neural networks. The complete list can be accessed [here](https://pytorch.org/docs/stable/nn.functional.html). Let's play with some activation functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(-5, 5, 100)\n",
    "\n",
    "plt.plot(x, F.relu(x), label='ReLU')\n",
    "plt.plot(x, F.elu(x), label='ELU')\n",
    "plt.plot(x, F.silu(x), label='SiLU')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### `nn.Module`\n",
    "\n",
    "PyTorch provides a very important base class named `nn.Module`. This class is used to build complex neural networks. In fact any class which inherits from it will automatically keep track of its parameters. To define your own `nn.Module` subclass you only need to implement the `__init__` constructor and the `forward` function. Let's see an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we create a class 'SimpleParametricModel' which inherits from nn.Module\n",
    "class SimpleParametricModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        # nn.Module's constructor has to be called BEFORE adding any sub-modules\n",
    "        super().__init__()\n",
    "        \n",
    "        # we add a linear layer to our module, which is itself a module\n",
    "        self.linear = nn.Linear(in_features=3, out_features=1, bias=True)\n",
    "        \n",
    "        # we add an activation function to our module\n",
    "        self.activation = nn.ReLU()\n",
    "        \n",
    "    # we define the forward pass\n",
    "    def forward(self, x):\n",
    "        x = self.linear(x)\n",
    "        x = self.activation(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create an instance of the `SimpleParametricModel` class and have a look at it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SimpleParametricModel()\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We observe that the two components are correctly registered. We can also inspect the parameters of our model with the `.parameters()` or `.state_dict()` methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for p in model.parameters():\n",
    "    print(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the parameters of the linear layer are considered as parameters of the model itself.\n",
    "\n",
    "PyTorch implements many components that are commonly used in neural networks such as convolutions, linear layers, normalizations and activations. Most of them are available as ready-to-use modules in `torch.nn` (imported as `nn` in this notebook).\n",
    "\n",
    "Let's build a simple multi-layer perceptron (MLP) with two hidden layers and a ReLU activation function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TwoLayerMLP(nn.Module):\n",
    "    def __init__(self, input_features, output_features, hidden_features=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.l1 = nn.Linear(input_features, hidden_features)\n",
    "        self.l2 = nn.Linear(hidden_features, hidden_features)\n",
    "        self.l3 = nn.Linear(hidden_features, output_features)\n",
    "        \n",
    "        self.activation = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        h1 = self.activation(self.l1(x))\n",
    "        h2 = self.activation(self.l2(h1))\n",
    "        y = self.l3(h2)\n",
    "        return y\n",
    "\n",
    "# instantiate the network to have 3 input features and 1 output feature (and 16 hidden features)\n",
    "network = TwoLayerMLP(3, 1)\n",
    "print(network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# we usually give batches of values to networks to take advantage of parallelism\n",
    "# here the batch size is 32\n",
    "x = torch.rand(8, 3)\n",
    "\n",
    "# call network.forward under the hood\n",
    "y = network(x)\n",
    "\n",
    "print(x, y, sep='\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### `nn.Sequential`\n",
    "\n",
    "Writing the forward pass can become cumbersome and dirty if you have a large number of sequential layers. For that kind of networks, you can use `nn.Sequential` as parent instead of `nn.Module`, which calls the layers sequentially within its forward pass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ElegantTwoLayerMLP(nn.Sequential):\n",
    "    def __init__(self, input_features, output_features, hidden_features=16):\n",
    "        super().__init__(\n",
    "            nn.Linear(input_features, hidden_features),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_features, hidden_features),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_features, output_features),\n",
    "        )\n",
    "\n",
    "#  instantiate the network to have 1 input feature and 1 output feature    \n",
    "network = ElegantTwoLayerMLP(1, 1)\n",
    "print(network)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Optimizers\n",
    "\n",
    "You should now be able to define any kind of neural network you like. However if you want to optimize it you would still need to iterate through all the parameters to perform a gradient descent step. Fortunately, `torch.nn.optim` (imported as `optim`) implements classes that handle that for you! For example, let's say you would like to learn the function $f(x) = x^2$ with the previous neural network, using a MSE loss (often called the criterion) and stochastic gradient descent (SGD):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return x**2\n",
    "\n",
    "# we use the MSE loss as criterion\n",
    "criterion = nn.MSELoss()\n",
    "\n",
    "# we create an instance of the SGD class that will make the updates for us\n",
    "optimizer = optim.SGD(params=network.parameters(), lr=.01)\n",
    "\n",
    "# Let's do some learning steps with randomly generated x values:\n",
    "for i in range(1000):\n",
    "    x = torch.randn(100, 1)\n",
    "    y = f(x)\n",
    "    y_pred = network(x)\n",
    "    \n",
    "    # we set all gradient values to 0\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    # we compute the loss\n",
    "    loss = criterion(y_pred, y)\n",
    "    \n",
    "    # we perform the backward pass to accumulate the gradients\n",
    "    loss.backward()\n",
    "    \n",
    "    # we update the parameters\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(-3, 3, 100).unsqueeze(1)\n",
    "\n",
    "with torch.no_grad():\n",
    "    y = network(x)\n",
    "\n",
    "plt.plot(x, f(x), label='$f(x)$')\n",
    "plt.plot(x, y, label='Approximation')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Training our first neural network classifier\n",
    "\n",
    "We now have all the necessary knowledge to build and train our very first neural network on a simple binary classification task. To do this we will start by importing a toy example dataset from the `sklearn` library, and then create and use the resulting splits for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.datasets import make_moons\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "X, Y = make_moons(500, noise=0.1)  # create artificial data\n",
    "\n",
    "X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.25, random_state=73)  # create splits\n",
    "\n",
    "plt.scatter(X_train[:,0], X_train[:,1], c=Y_train)  # visualize the training data\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When we train a neural network in PyTorch, no matter how complex the model is, we always go through a training loop. In this loop we feed the data to the model and get its predictions. We then compare the predictions of the network to the ground truth and adjust the parameters of the model by performing gradient descent. We have already seen all the components that are necessary for going through this process, so the only thing that remains to be done is to put all the pieces of this notebook together and train our **first awesome neural network!**\n",
    "\n",
    "During the training stage we would like to keep track whether our model will improve over the different iterations. It is therefore good practice to monitor whether the loss we are minimizing decreases over time, and whether the overall performance of the model increases the more training iterations we perform. Remember that one of the key components of PyTorch are tensors, and we might not always have the data coming in this specific format. It is therefore necessary to convert it if needed.\n",
    "\n",
    "Once the data is in an appropriate format it can be given to the model, and we can obtain its predictions. This is what we usually call the **forward pass**. Once we obtain our predictions, we can compare how close they are to what we would like the network to predict: to do this we feed our predictions together with the true labels through the loss function we are minimizing. At its early training stages the network will perform poorly, but it will improve as its weights/parameters are updated by gradient descent. We can do this very easily by obtaining the gradients of the parameters with respect to the loss function we are minimizing (**backward pass**) and adjusting these weights with the optimizer we defined previously.\n",
    "\n",
    "Once all of this is done we can measure the performance of our model, which in this case will be reflected by its accuracy in classifying the synthetic dataset we created beforehand. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = nn.Sequential(\n",
    "    nn.Linear(X.shape[-1], 32),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(32, 32),\n",
    "    nn.ReLU(),\n",
    "    nn.Linear(32, 1),\n",
    "    nn.Sigmoid(),\n",
    ")\n",
    "criterion = nn.BCELoss()\n",
    "optimizer = optim.Adam(net.parameters(), lr=.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loss = []  # where we keep track of the loss\n",
    "train_accuracy = []  # where we keep track of the accuracy of the model\n",
    "iters = 128  # number of training iterations\n",
    "\n",
    "X_train_t = torch.from_numpy(X_train).float()\n",
    "Y_train_t = torch.from_numpy(Y_train).float()\n",
    "\n",
    "for i in range(iters):\n",
    "    y_pred = net(X_train_t)  # forward pass\n",
    "    y_pred = y_pred.squeeze(-1)  # transform the 1-element vectors into scalars\n",
    "\n",
    "    optimizer.zero_grad() # reset the gradients to 0\n",
    "    loss = criterion(y_pred, Y_train_t)  # compute the loss\n",
    "    loss.backward()  # obtain the gradients with respect to the loss\n",
    "    optimizer.step()  # perform one step of gradient descent\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        y_pred_class = y_pred > 0.5  # we assign a label (0 or 1) based on the network's prediction\n",
    "        accuracy = (Y_train_t == y_pred_class).float().mean()  # compute accuracy\n",
    "    \n",
    "    train_accuracy.append(accuracy.item())\n",
    "    train_loss.append(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 1)\n",
    "\n",
    "axs[0].set_title('Training Loss')\n",
    "axs[0].plot(train_loss)\n",
    "axs[0].set_xlabel('Iterations')\n",
    "axs[0].set_ylabel('Loss (Binary Cross Entropy)')\n",
    "\n",
    "axs[1].set_title('Training Accuracy')\n",
    "axs[1].plot(train_accuracy)\n",
    "axs[1].set_xlabel('Iterations')\n",
    "axs[1].set_ylabel('Accuracy')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_decision_boundary(X, y, model, steps=1000, cmap='Paired'):\n",
    "    \"\"\"\n",
    "    Function to plot the decision boundary and data points of a model.\n",
    "    Data points are colored based on their actual label.\n",
    "    \"\"\"\n",
    "    cmap = plt.get_cmap(cmap)\n",
    "\n",
    "    # Define region of interest by data limits\n",
    "    xmin, xmax = X[:,0].min(), X[:,0].max()\n",
    "    ymin, ymax = X[:,1].min(), X[:,1].max()\n",
    "    steps = 1000\n",
    "    x_span = np.linspace(xmin - .5, xmax + .5, steps)\n",
    "    y_span = np.linspace(ymin - .5, ymax + .5, steps)\n",
    "    xx, yy = np.meshgrid(x_span, y_span)\n",
    "\n",
    "    # Make predictions across region of interest\n",
    "    with torch.no_grad():\n",
    "        grid = np.stack([xx, yy], axis=-1)\n",
    "        grid = torch.from_numpy(grid).float()\n",
    "        pred = model(grid).squeeze(-1)\n",
    "\n",
    "    # Plot decision boundary in region of interest\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.contourf(xx, yy, pred, cmap=cmap, alpha=0.5)\n",
    "\n",
    "    # Get predicted labels on training data and plot\n",
    "    train_labels = model(torch.from_numpy(X).float())\n",
    "    ax.scatter(X[:,0], X[:,1], c=y.ravel(), cmap=cmap, lw=0)\n",
    "\n",
    "    return fig\n",
    "\n",
    "fig = plot_decision_boundary(X_test, Y_test, net, cmap = 'RdBu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Discuss the bugs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After actively following the theoretical lectures of professor Louppe, our favourite Italian student Al Dente remembered that a common loss function one can use when dealing with classification problems is the Cross-Entropy loss. However after implementing it in PyTorch, Al was disappointed because his code was not working.\n",
    "\n",
    "What do you think are the causes of the bug(s)? Maybe, first check the documentation of [`torch.nn.CrossEntropyLoss`](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html). Briefly discuss your ideas and **bonus** can you come up with a solution?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = nn.Sequential(\n",
    "    nn.Linear(X.shape[-1], 32),\n",
    "    nn.ELU(),\n",
    "    nn.Linear(32, 32),\n",
    "    nn.ELU(),\n",
    "    nn.Linear(32, 1),\n",
    ")\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=0.01)\n",
    "\n",
    "train_loss = []\n",
    "train_accuracy = []\n",
    "iters = 420\n",
    "\n",
    "X_train_t = torch.from_numpy(X_train).float()\n",
    "Y_train_t = torch.from_numpy(Y_train).float()\n",
    "\n",
    "for i in range(iters):\n",
    "    y_pred = net(X_train_t)\n",
    "    y_pred = y_pred.squeeze(-1)\n",
    "\n",
    "    loss = criterion(y_pred, Y_train_t)\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> your discussion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code (bonus)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Try it yourself!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now try to play around with the neural network yourself: you are free to experiment with whatever you think is best exploring. Your goal is to come up with three modifications to the provided network architecture and to compare the performance of your modified models. As potential modifications you could experiment with changing the depth or width of the model, modifying the activation function, learning rate, optimizer, etc.\n",
    "\n",
    "Report your results with some appropriate plots and with a brief presentation of what you observe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> your discussion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feedback\n",
    "\n",
    "We will now ask you a few questions to improve the content of this homework for next years."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:blue\">How long did you spend on this homework?</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:blue\">Did you learn something?</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:blue\">Do you now feel comfortable with writing simple mathematical operations (like you would do in matlab or with numpy) in PyTorch?</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:blue\">Do you now feel comfortable using PyTorch for performing gradient descent?</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span style=\"color:blue\">Do you now feel comfortable with building and learning neural networks with PyTorch?</span>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
